"""
    $(TYPEDSIGNATURES)

Given a function expression `expr`, return a callable version of it.

# Keyword arguments
- `eval_expression`: Whether to use `eval` to make `expr` callable. If `false`, uses
  RuntimeGeneratedFunctions.jl.
- `eval_module`: The module to `eval` the expression `expr` in. If `!eval_expression`,
  this is the cache and context module for the `RuntimeGeneratedFunction`.
"""
function eval_or_rgf(expr::Expr; eval_expression = false, eval_module = @__MODULE__)
    if eval_expression
        return eval_module.eval(expr)
    else
        return drop_expr(RuntimeGeneratedFunction(eval_module, eval_module, expr))
    end
end

eval_or_rgf(::Nothing; kws...) = nothing

"""
    $(TYPEDSIGNATURES)

Return the name for the `i`th argument in a function generated by `build_function_wrapper`.
"""
function generated_argument_name(i::Int)
    return Symbol(:__mtk_arg_, i)
end

"""
    $(TYPEDSIGNATURES)

Given the arguments to `build_function_wrapper`, return a list of assignments which
reconstruct array variables if they are present scalarized in `args`.

# Keyword Arguments

- `argument_name` a function of the form `(::Int) -> Symbol` which takes the index of
  an argument to the generated function and returns the name of the argument in the
  generated function.
"""
function array_variable_assignments(args...; argument_name = generated_argument_name)
    # map array symbolic to an identically sized array where each element is (buffer_idx, idx_in_buffer)
    var_to_arridxs = Dict{BasicSymbolic, Array{Tuple{Int, Int}}}()
    for (i, arg) in enumerate(args)
        # filter out non-arrays
        # any element of args which is not an array is assumed to not contain a
        # scalarized array symbolic. This works because the only non-array element
        # is the independent variable
        symbolic_type(arg) == NotSymbolic() || continue
        arg isa AbstractArray || continue

        # go through symbolics
        for (j, var) in enumerate(arg)
            var = unwrap(var)
            # filter out non-array-symbolics
            iscall(var) || continue
            operation(var) == getindex || continue
            arrvar = arguments(var)[1]
            # get and/or construct the buffer storing indexes
            idxbuffer = get!(
                () -> map(Returns((0, 0)), eachindex(arrvar)), var_to_arridxs, arrvar)
            Origin(first.(axes(arrvar))...)(idxbuffer)[arguments(var)[2:end]...] = (i, j)
        end
    end

    assignments = Assignment[]
    for (arrvar, idxs) in var_to_arridxs
        # all elements of the array need to be present in `args` to form the
        # reconstructing assignment
        any(iszero ∘ first, idxs) && continue

        # if they are all in the same buffer, we can take a shortcut and `view` into it
        if allequal(Iterators.map(first, idxs))
            buffer_idx = first(first(idxs))
            idxs = map(last, idxs)
            # if all the elements are contiguous and ordered, turn the array of indexes into a range
            # to help reduce allocations
            if first(idxs) < last(idxs) && vec(idxs) == first(idxs):last(idxs)
                idxs = first(idxs):last(idxs)
            elseif vec(idxs) == first(idxs):-1:last(idxs)
                idxs = first(idxs):-1:last(idxs)
            else
                # Otherwise, turn the indexes into an `SArray` so they're stack-allocated
                idxs = SArray{Tuple{size(idxs)...}}(idxs)
            end
            # view and reshape

            expr = term(reshape, term(view, argument_name(buffer_idx), idxs),
                size(arrvar))
        else
            elems = map(idxs) do idx
                i, j = idx
                term(getindex, argument_name(i), j; type = Real, shape = SU.ShapeVecT())
            end
            # use `MakeArray` syntax and generate a stack-allocated array
            expr = term(SymbolicUtils.Code.create_array, SArray, nothing,
                Val(ndims(arrvar)), Val(length(arrvar)), elems...)
        end
        if any(x -> !isone(first(x)), axes(arrvar))
            expr = term(Origin(first.(axes(arrvar))...), expr)
        end
        push!(assignments, arrvar ← expr)
    end

    return assignments
end

"""
    $(TYPEDSIGNATURES)

Check if the variable `var` is a delayed variable, where `iv` is the independent
variable.
"""
function isdelay(var, iv)
    iv === nothing && return false
    if iscall(var) && ModelingToolkitBase.isoperator(var, Differential)
        return isdelay(arguments(var)[1], iv)
    end
    isvariable(var) || return false
    isparameter(var) && return false
    if iscall(var) && !ModelingToolkitBase.isoperator(var, Symbolics.Operator)
        args = arguments(var)
        length(args) == 1 || return false
        arg = args[1]
        isequal(arg, iv) && return false
        iscall(arg) || return true
        issym(operation(arg)) && !iscalledparameter(arg) && return false
        return true
    end
    return false
end

"""
The argument of generated functions corresponding to the history function.
"""
const DDE_HISTORY_FUN = SSym(:___history___; type = SU.FnType{Tuple{Any, <:Real}, Vector{Real}, Nothing}, shape = SU.Unknown(1))
const BVP_SOLUTION = SSym(:__sol__; type = Symbolics.FnType{Tuple{<:Real}, Vector{Real}, Nothing}, shape = SU.Unknown(1))

"""
    $(TYPEDSIGNATURES)

Turn delayed unknowns in `eqs` into calls to `DDE_HISTORY_FUNCTION`.

# Arguments

- `sys`: The system of DDEs.
- `eqs`: The equations to convert.

# Keyword Arguments

- `param_arg`: The name of the variable containing the parameter object.
- `histfn`: The history function to use for codegen, called as `histfn(p, t)`
"""
function delay_to_function(
        sys::AbstractSystem, eqs = full_equations(sys); param_arg = MTKPARAMETERS_ARG, histfn = DDE_HISTORY_FUN)
    delay_to_function(eqs,
        get_iv(sys),
        Dict{Any, Int}(operation(s) => i for (i, s) in enumerate(unknowns(sys))),
        parameters(sys),
        histfn; param_arg)
end
function delay_to_function(eqs::Vector, iv, sts, ps, h; param_arg = MTKPARAMETERS_ARG)
    delay_to_function.(eqs, (iv,), (sts,), (ps,), (h,); param_arg)
end
function delay_to_function(eq::Equation, iv, sts, ps, h; param_arg = MTKPARAMETERS_ARG)
    delay_to_function(eq.lhs, iv, sts, ps, h; param_arg) ~ delay_to_function(
        eq.rhs, iv, sts, ps, h; param_arg)
end
function delay_to_function(expr, iv, sts, ps, h; param_arg = MTKPARAMETERS_ARG)
    if isdelay(expr, iv)
        v = operation(expr)
        time = arguments(expr)[1]
        idx = sts[v]
        return term(getindex, h(param_arg, time), idx, type = Real)
    elseif iscall(expr)
        return maketerm(typeof(expr),
            operation(expr),
            map(x -> delay_to_function(x, iv, sts, ps, h; param_arg), arguments(expr)),
            metadata(expr))
    else
        return expr
    end
end

function __search_dervars_recurse(x::SymbolicT)
    iscall(x) && Moshi.Match.@match x begin
        BSImpl.Term(; f) && if f isa Operator end => false
        _ => true
    end
end

"""
    $(TYPEDSIGNATURES)

A wrapper around `build_function` which performs the necessary transformations for
code generation of all types of systems. `expr` is the expression returned from the
generated functions, and `args` are the arguments.

# Keyword Arguments

- `p_start`, `p_end`: Denotes the indexes in `args` where the buffers of the splatted
  `MTKParameters` object are present. These are collapsed into a single argument and
  destructured inside the function. `p_start` must also be provided for non-split systems
  since it is used by `wrap_delays`.
- `wrap_delays`: Whether to transform delayed unknowns of `sys` present in `expr` into
  calls to a history function. The history function is added to the list of arguments
  right before parameters, at the index `p_start`.
- `histfn`: The history function to use for transforming delayed terms. For any delayed
  term `x(expr)`, this is called as `histfn(p, expr)` where `p` is the parameter object.
- `histfn_symbolic`: The symbolic history function variable to add as an argument to the
  generated function.
- `wrap_code`: Forwarded to `build_function`.
- `add_observed`: Whether to add assignment statements for observed equations in the
  generated code.
- `filter_observed`: A predicate function to filter out observed equations which should
  not be added to the generated code.
- `create_bindings`: Whether to explicitly destructure arrays of symbolics present in
  `args` in the generated code. If `false`, all usages of the individual symbolics will
  instead call `getindex` on the relevant argument. This is useful if the generated
  function writes to one of its arguments and expects subsequent code to use the new
  values. Note that the collapsed `MTKParameters` argument will always be explicitly
  destructured regardless of this keyword argument.
- `output_type`: The type of the output buffer. If `mkarray` (see below) is `nothing`,
  this will be passed to the `similarto` argument of `build_function`. If `output_type`
  is `Tuple`, `expr` will be wrapped in `SymbolicUtils.Code.MakeTuple` (regardless of
  whether it is scalar or an array).
- `mkarray`: A function which accepts `expr` and `output_type` and returns a code
  generation object similar to `MakeArray` or `MakeTuple` to be used to generate
  code for `expr`.
- `wrap_mtkparameters`: Whether to collapse parameter buffers for a split system into a
  argument.
- `extra_assignments`: Extra `Assignment` statements to prefix to `expr`, after all other
  assignments.

All other keyword arguments are forwarded to `build_function`.
"""
function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
        p_end = is_time_dependent(sys) ? length(args) - 1 : length(args),
        wrap_delays = is_dde(sys), histfn = DDE_HISTORY_FUN, histfn_symbolic = histfn, wrap_code = identity,
        add_observed = true, filter_observed = Returns(true),
        create_bindings = false, output_type = nothing, mkarray = nothing,
        wrap_mtkparameters = true, extra_assignments = Assignment[], cse = true, kwargs...)
    isscalar = !(expr isa AbstractArray || symbolic_type(expr) == ArraySymbolic())
    # filter observed equations
    obs = filter(filter_observed, observed(sys))
    # turn delayed unknowns into calls to the history function
    if wrap_delays
        param_arg = is_split(sys) ? MTKPARAMETERS_ARG : generated_argument_name(p_start)
        obs = map(obs) do eq
            delay_to_function(sys, eq; param_arg, histfn)
        end
        expr = delay_to_function(sys, expr; param_arg, histfn)
        # add extra argument
        args = (args[1:(p_start - 1)]..., histfn_symbolic, args[p_start:end]...)
        p_start += 1
        p_end += 1
    end

    dervars = Dict{SymbolicT, SymbolicT}()
    dervars_in_expr = Set{SymbolicT}()
    if isscheduled(sys)
        sched::Schedule = get_schedule(sys)
        for (k, v) in sched.dummy_sub
            ttk = default_toterm(k)
            isequal(ttk, v) && continue
            dervars[default_toterm(k)] = v
        end
    else
        for eq in equations(sys)
            isdiffeq(eq) || continue
            ttk = default_toterm(eq.lhs)
            isequal(ttk, eq.rhs) && continue
            dervars[ttk] = eq.rhs
        end
    end
    Symbolics.get_variables!(dervars_in_expr, expr, keys(dervars); recurse = __search_dervars_recurse)
    # only get the necessary observed equations, avoiding extra computation
    if add_observed && !isempty(obs)
        obsidxs = BitSet(observed_equations_used_by(sys, expr; obs))
    else
        obsidxs = BitSet()
    end
    # similarly for parameter dependency equations
    reqd_bound_pars = OrderedSet{SymbolicT}()
    bgraph::ParameterBindingsGraph = if iscomplete(sys)
        get_parameter_bindings_graph(sys)
    else
        ParameterBindingsGraph(sys)
    end
    
    bound_parameters_used_by!(reqd_bound_pars, sys, expr; bgraph)
    for i in obsidxs
        bound_parameters_used_by!(reqd_bound_pars, sys, obs[i].rhs; bgraph)
    end
    for dervar in dervars_in_expr
        bound_parameters_used_by!(reqd_bound_pars, sys, dervars[dervar]; bgraph)
        union!(obsidxs, observed_equations_used_by(sys, dervars[dervar]; obs))
    end
    sort_bound_parameters!(reqd_bound_pars, sys; bgraph)
    # assignments for reconstructing scalarized array symbolics
    assignments = array_variable_assignments(args...)
    binds = bindings(sys)
    for p in reqd_bound_pars
        push!(assignments, p ← binds[p])
    end
    obsidxs = collect(obsidxs)
    for eq in obs[obsidxs]
        push!(assignments, eq.lhs ← eq.rhs)
    end

    for dervar in dervars_in_expr
        push!(assignments, dervar ← dervars[dervar])
    end
    append!(assignments, extra_assignments)

    args = ntuple(Val(length(args))) do i
        arg = args[i]
        # Make sure to use the proper names for arguments
        if symbolic_type(arg) == NotSymbolic() && arg isa AbstractArray
            DestructuredArgs(arg, generated_argument_name(i); create_bindings)
        else
            arg
        end
    end

    # wrap into a single MTKParameters argument
    if is_split(sys) && wrap_mtkparameters
        if p_start > p_end
            # In case there are no parameter buffers, still insert an argument
            args = (args[1:(p_start - 1)]..., MTKPARAMETERS_ARG, args[(p_end + 1):end]...)
        else
            # cannot apply `create_bindings` here since it doesn't nest
            args = (args[1:(p_start - 1)]...,
                DestructuredArgs(collect(args[p_start:p_end]), MTKPARAMETERS_ARG),
                args[(p_end + 1):end]...)
        end
    end

    # add preface assignments
    if has_preface(sys) && (pref = preface(sys)) !== nothing
        append!(assignments, pref)
    end

    wrap_code = wrap_code .∘ wrap_assignments(isscalar, assignments)

    # handling of `output_type` and `mkarray`
    similarto = nothing
    if output_type === Tuple
        expr = MakeTuple(Tuple(expr))
        wrap_code = wrap_code[1]
    elseif mkarray === nothing
        similarto = output_type
    else
        expr = mkarray(expr, output_type)
        wrap_code = wrap_code[2]
    end

    # scalar `build_function` only accepts a single function for `wrap_code`.
    if wrap_code isa Tuple && symbolic_type(expr) == ScalarSymbolic()
        wrap_code = wrap_code[1]
    end
    return build_function(expr, args...; wrap_code, similarto, cse, kwargs...)
end

"""
    $(TYPEDEF)

A wrapper around a generated in-place and out-of-place function. The type-parameter `P`
must be a 3-tuple where the first element is the index of the parameter object in the
arguments, the second is the expected number of arguments in the out-of-place variant
of the function, and the third is a boolean indicating whether the generated functions
are for a split system. For scalar functions, the inplace variant can be `nothing`.
"""
struct GeneratedFunctionWrapper{P, O, I} <: Function
    f_oop::O
    f_iip::I
end

function GeneratedFunctionWrapper{P}(foop::O, fiip::I) where {P, O, I}
    GeneratedFunctionWrapper{P, O, I}(foop, fiip)
end

function GeneratedFunctionWrapper{P}(::Type{Val{true}}, foop, fiip; kwargs...) where {P}
    :($(GeneratedFunctionWrapper{P})($foop, $fiip))
end

function GeneratedFunctionWrapper{P}(::Type{Val{false}}, foop, fiip; kws...) where {P}
    GeneratedFunctionWrapper{P}(eval_or_rgf(foop; kws...), eval_or_rgf(fiip; kws...))
end

function (gfw::GeneratedFunctionWrapper)(args...)
    _generated_call(gfw, args...)
end

@generated function _generated_call(gfw::GeneratedFunctionWrapper{P}, args...) where {P}
    paramidx, nargs, issplit = P
    iip = false
    # IIP case has one more argument
    if length(args) == nargs + 1
        nargs += 1
        paramidx += 1
        iip = true
    end
    if length(args) != nargs
        throw(ArgumentError("Expected $nargs arguments, got $(length(args))."))
    end

    # the function to use
    f = iip ? :(gfw.f_iip) : :(gfw.f_oop)
    # non-split systems just call it as-is
    if !issplit
        return :($f(args...))
    end
    if args[paramidx] <: Union{Tuple, MTKParameters} &&
       !(args[paramidx] <: Tuple{Vararg{Number}})
        # for split systems, call it as-is if the parameter object is a tuple or MTKParameters
        # but not if it is a tuple of numbers
        return :($f(args...))
    else
        # The user provided a single buffer/tuple for the parameter object, so wrap that
        # one in a tuple
        fargs = ntuple(Val(length(args))) do i
            i == paramidx ? :((args[$i], nothing)) : :(args[$i])
        end
        return :($f($(fargs...)))
    end
end

"""
    $(TYPEDSIGNATURES)

Optionally compile a method and optionally wrap it in a `GeneratedFunctionWrapper` on the
basis of `expression` `wrap_gfw`, both of type `Union{Type{Val{true}}, Type{Val{false}}}`.
`gfw_args` is the first type parameter of `GeneratedFunctionWrapper`. `f` is a tuple of
function expressions of the form `(oop, iip)` or a single out-of-place function expression.
Keyword arguments are forwarded to `eval_or_rgf`.
"""
function maybe_compile_function(expression, wrap_gfw::Type{Val{true}},
        gfw_args::Tuple{Int, Int, Bool}, f::NTuple{2, Expr}; kwargs...)
    GeneratedFunctionWrapper{gfw_args}(expression, f...; kwargs...)
end

function maybe_compile_function(expression::Type{Val{false}}, wrap_gfw::Type{Val{false}},
        gfw_args::Tuple{Int, Int, Bool}, f::NTuple{2, Expr}; kwargs...)
    eval_or_rgf.(f; kwargs...)
end

function maybe_compile_function(expression::Type{Val{true}}, wrap_gfw::Type{Val{false}},
        gfw_args::Tuple{Int, Int, Bool}, f::Union{Expr, NTuple{2, Expr}}; kwargs...)
    return f
end

function maybe_compile_function(expression, wrap_gfw::Type{Val{true}},
        gfw_args::Tuple{Int, Int, Bool}, f::Expr; kwargs...)
    GeneratedFunctionWrapper{gfw_args}(expression, f, nothing; kwargs...)
end

function maybe_compile_function(expression::Type{Val{false}}, wrap_gfw::Type{Val{false}},
        gfw_args::Tuple{Int, Int, Bool}, f::Expr; kwargs...)
    eval_or_rgf(f; kwargs...)
end
